import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import time

# 简单的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 8, kernel_size=3, padding=1)  # 减少通道数
        self.conv2 = nn.Conv2d(8, 16, kernel_size=3, padding=1)  # 减少通道数
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(16 * 7 * 7, 32)  # 减少神经元数量
        self.fc2 = nn.Linear(32, 10)
        self.relu = nn.ReLU()
        
    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 16 * 7 * 7)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

def test_mps_completely():
    """彻底测试MPS是否可用"""
    try:
        # 检查基本属性
        if not hasattr(torch.backends, 'mps'):
            return False
            
        if not torch.backends.mps.is_available():
            return False
            
        if not torch.backends.mps.is_built():
            return False
        
        # 实际测试MPS功能
        device = torch.device('mps')
        
        # 测试张量创建
        a = torch.randn(10, 10).to(device)
        b = torch.randn(10, 10).to(device)
        
        # 测试基本运算
        c = torch.mm(a, b)
        
        # 测试反向传播
        d = c.sum()
        d.backward()
        
        # 如果所有测试都通过，MPS可用
        return True
        
    except Exception as e:
        print(f"MPS测试失败: {e}")
        return False

def get_device():
    """获取可用设备"""
    print("检测设备...")
    
    # 首先彻底测试MPS
    if test_mps_completely():
        device = torch.device('mps')
        print("MPS设备通过完整测试，可以使用")
    elif torch.cuda.is_available():
        device = torch.device('cuda')
        print("使用CUDA设备")
    else:
        device = torch.device('cpu')
        print("使用CPU设备")
        
    return device

def main():
    print("保守型CNN示例 (MPS安全版本)")
    print("=" * 40)
    
    # 获取设备
    device = get_device()
    print(f"使用的设备: {device}")
    
    # 数据预处理
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    # 加载数据（小批量）
    print("加载数据...")
    trainset = torchvision.datasets.MNIST(
        root='./mnist_data', 
        train=True, 
        download=True, 
        transform=transform
    )
    
    trainloader = torch.utils.data.DataLoader(
        trainset, 
        batch_size=16,  # 更小的批次
        shuffle=True,
        num_workers=0
    )
    
    # 创建模型
    print("创建模型...")
    model = SimpleCNN()
    
    try:
        model.to(device)
        print(f"模型成功迁移到 {device} 设备")
    except Exception as e:
        print(f"模型迁移失败: {e}")
        print("使用CPU设备继续...")
        device = torch.device('cpu')
        model.to(device)
        
    print(f"模型参数数量: {sum(p.numel() for p in model.parameters())}")
    
    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    
    # 训练几个批次来测试
    print("开始训练测试...")
    model.train()
    start_time = time.time()
    
    successful_batches = 0
    max_test_batches = 50  # 只测试少量批次
    
    for i, (inputs, labels) in enumerate(trainloader):
        try:
            # 数据迁移到设备
            inputs, labels = inputs.to(device), labels.to(device)
            
            # 前向传播
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            successful_batches += 1
            if successful_batches % 10 == 0:
                print(f'成功完成 {successful_batches} 个批次，当前损失: {loss.item():.4f}')
                
            # 达到测试数量后停止
            if successful_batches >= max_test_batches:
                break
                
        except Exception as e:
            print(f"批次 {i} 出错: {e}")
            print("跳过当前批次...")
            continue
    
    training_time = time.time() - start_time
    print(f"训练测试完成，成功处理 {successful_batches} 个批次，耗时: {training_time:.2f}秒")
    
    # 简单评估
    print("进行评估测试...")
    model.eval()
    correct = 0
    total = 0
    
    eval_batches = 0
    max_eval_batches = 20
    
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(trainloader):
            try:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                
                eval_batches += 1
                if eval_batches >= max_eval_batches:
                    break
                    
            except Exception as e:
                print(f"评估批次 {i} 出错: {e}")
                continue
    
    if total > 0:
        accuracy = 100 * correct / total
        print(f'评估准确率: {accuracy:.2f}%')
    
    print("程序执行完成!")

if __name__ == "__main__":
    main()